import gradio as gr
import time
import json
import re
import requests
from pypinyin import lazy_pinyin, Style


# 拼音转换工具类
class PinyinConverter:
    @staticmethod
    def get_pinyin(text):
        """将中文文本转换为带声调的拼音，忽略非中文字符"""
        if not text or not isinstance(text, str):
            return ""
        # 只保留中文字符进行拼音转换
        chinese_chars = [char for char in text if '\u4e00' <= char <= '\u9fff']
        if not chinese_chars:
            return ""
        # 转换为带声调的拼音，空格分隔
        pinyin_list = lazy_pinyin(''.join(chinese_chars), style=Style.TONE)
        return ' '.join(pinyin_list)

    @staticmethod
    def extract_and_add_pinyin(content):
        """提取内容中【】包裹的诗句，并为其添加拼音标注"""
        if not content:
            return content

        # 正则表达式匹配【】包裹的内容
        pattern = r'【([^】]+)】'

        def replace_with_pinyin(match):
            poem_line = match.group(1)
            pinyin = PinyinConverter.get_pinyin(poem_line)
            if pinyin:
                return f'【{poem_line}】{pinyin}'
            else:
                return match.group(0)

        processed_content = re.sub(pattern, replace_with_pinyin, content)
        return processed_content


# 诗歌API调用类
class PoetryAPI:
    def __init__(self, api_key, base_url, model="ernie-bot-turbo", debug=False, rpm_limit=10):
        self.api_key = api_key
        self.base_url = base_url.rstrip('/')
        self.model = model
        self.debug = debug
        self.rpm_limit = rpm_limit
        self.min_request_interval = 2.0
        self.last_request_time = 0
        self.request_timestamps = []
        self.is_cancelled = False

    def handle_rate_limit(self):
        current_time = time.time()

        if self.is_cancelled:
            return False

        # 检查最小请求间隔
        if current_time - self.last_request_time < self.min_request_interval:
            wait_time = self.min_request_interval - (current_time - self.last_request_time)
            time.sleep(wait_time)
            current_time = time.time()

        # 清理1分钟前的请求记录
        self.request_timestamps = [t for t in self.request_timestamps
                                   if current_time - t < 60]

        # 检查RPM限制
        if len(self.request_timestamps) >= self.rpm_limit:
            wait_time = 60 - (current_time - self.request_timestamps[0]) + 1
            time.sleep(wait_time)
            current_time = time.time()
            self.request_timestamps = [t for t in self.request_timestamps
                                       if current_time - t < 60]

        # 记录当前请求时间
        self.request_timestamps.append(current_time)
        self.last_request_time = current_time

        if len(self.request_timestamps) > self.rpm_limit * 2:
            self.request_timestamps = self.request_timestamps[-self.rpm_limit:]

        return True

    def couplet_stream(self, upper_sentence, user_requirements=""):
        try:
            self.is_cancelled = False

            if not self.handle_rate_limit():
                return "", "请求频率超限，取消本次请求"

            # 构建请求参数
            url = f"{self.base_url}/chat/completions"
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.api_key}"
            }

            base_prompt = f"""你是专为小朋友服务的诗歌助手，现在要和1-6年级的孩子玩对诗游戏。本次对话独立，不参考历史记录。

【核心规则 - 必须严格遵守】：
1. 对于输入的诗句上半句"{upper_sentence}"：
   - 如果是经典诗句（如唐诗宋词），必须用原句的下半句，不能修改
   - 如果不是经典诗句，根据语义和韵律创作适合小朋友的下半句
2. 回应必须符合儿童理解能力，语言*极度简单*、口语化、生动，避免复杂词汇。

【回应结构 - 按顺序输出】：
1. 先给出下半句，用【】包裹（例如：【疑是地上霜】）
2. 用*幼儿园小朋友能听懂的话*解释两句诗的意思，每句解释不超过20个字，要非常直白。
3. 说明诗句是否为经典名句：
   - 若是，注明作者和诗名（简洁，如：李白《静夜思》）
   - 若不是，说明这是为小朋友创作的
4. 最后用一句鼓励的话，激发孩子对古诗的兴趣

【特别说明】：本次回应仅针对"{upper_sentence}"，不涉及其他内容。"""

            if user_requirements.strip():
                base_prompt += f"\n\n【用户额外要求】：{user_requirements}"

            payload = {
                "model": self.model,
                "temperature": 0.3,
                "messages": [{"role": "user", "content": base_prompt}],
                "stream": True
            }

            # 发送POST请求
            response = requests.post(
                url,
                headers=headers,
                json=payload,
                stream=True
            )
            response.raise_for_status()

            result = ""
            # 处理流式响应
            for line in response.iter_lines():
                if self.is_cancelled:
                    return result, "请求已取消"

                if line:
                    line = line.decode('utf-8').lstrip('data: ').strip()
                    if line == "[DONE]":
                        break

                    try:
                        chunk = json.loads(line)
                        if chunk.get('choices') and len(chunk['choices']) > 0:
                            chunk_content = chunk['choices'][0].get('delta', {}).get('content', '')
                            if chunk_content:
                                processed_chunk = PinyinConverter.extract_and_add_pinyin(chunk_content)
                                result += processed_chunk
                    except json.JSONDecodeError:
                        continue
                time.sleep(0.05)

            return result, "生成成功"

        except Exception as e:
            error_str = f"对句请求失败: {str(e)}"
            if "invalid_model" in str(e):
                error_str += "\n提示：模型不存在或无权限"
            elif "401" in str(e):
                error_str += "\n提示：API Key无效，请检查密钥"
            elif "429" in str(e):
                error_str += "\n提示：请求过于频繁，请稍后再试"
            return "", error_str

    def generate_poem_stream(self, keywords, style, user_requirements=""):
        try:
            self.is_cancelled = False

            if not keywords:
                return "", "关键词不能为空"
            if not style:
                return "", "请选择诗歌风格"

            if not self.handle_rate_limit():
                return "", "请求频率超限，取消本次请求"

            style_guide = {
                "唐诗": "句式整齐，多为五言或七言，押韵对仗工整，意境深远",
                "宋词": "句子长短不一，注重情感表达，风格多样（豪放/婉约）",
                "五言绝句": "四句，每句五个字，结构紧凑，前两句写景后两句抒情",
                "七言绝句": "四句，每句七个字，节奏明快，适合表达瞬间感悟"
            }

            base_prompt = f"""你是儿童诗歌创作家，要为1-6年级小朋友创作{style}。本次创作独立，不参考历史记录。

【创作规则】：
1. 语言必须简单易懂，使用小朋友熟悉的词汇
2. 所有诗句都要标注拼音
3. 若引用经典诗句，必须使用原句并注明出处

【创作流程】：
1. 先推荐2首与关键词"{keywords}"相关的经典古诗，每首包含：
   - 完整诗句（每句用【】包裹）
   - 每句的拼音标注
   - 简单解释（小朋友能懂）
   - 作者和诗名
2. 然后基于关键词创作符合{style}特点的原创诗句：{style_guide[style]}
3. 原创诗句必须包含所有关键词：{keywords}
4. 如果找不到相关的经典古诗，就直接创作原创诗句

【回应结构】：
1. 经典古诗推荐部分（如有）
2. 原创{style}部分，每句诗单独一行并用【】包裹，带拼音
3. 逐句解释原创诗句意思，用*幼儿园小朋友能听懂的话*，每句解释不超过20个字
4. 说明原创诗中描绘的画面和表达的情感，语言要简单

【特别说明】：本次创作仅针对关键词{keywords}，不涉及其他内容。"""

            if user_requirements.strip():
                base_prompt += f"\n\n【用户额外要求】：{user_requirements}"

            # 构建请求参数
            url = f"{self.base_url}/chat/completions"
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.api_key}"
            }
            payload = {
                "model": self.model,
                "temperature": 0.7,
                "messages": [{"role": "user", "content": base_prompt}],
                "stream": True
            }

            # 发送POST请求
            response = requests.post(
                url,
                headers=headers,
                json=payload,
                stream=True
            )
            response.raise_for_status()

            result = ""
            # 处理流式响应
            for line in response.iter_lines():
                if self.is_cancelled:
                    return result, "请求已取消"

                if line:
                    line = line.decode('utf-8').lstrip('data: ').strip()
                    if line == "[DONE]":
                        break

                    try:
                        chunk = json.loads(line)
                        if chunk.get('choices') and len(chunk['choices']) > 0:
                            chunk_content = chunk['choices'][0].get('delta', {}).get('content', '')
                            if chunk_content:
                                processed_chunk = PinyinConverter.extract_and_add_pinyin(chunk_content)
                                result += processed_chunk
                    except json.JSONDecodeError:
                        continue
                time.sleep(0.05)

            return result, "生成成功"

        except Exception as e:
            error_str = f"生成诗句失败: {str(e)}"
            if "invalid_model" in str(e):
                error_str += "\n提示：模型不存在或无权限"
            elif "401" in str(e):
                error_str += "\n提示：API Key无效，请检查密钥"
            elif "429" in str(e):
                error_str += "\n提示：请求过于频繁，请稍后再试"
            return "", error_str

    def cancel_request(self):
        self.is_cancelled = True


# Gradio 界面逻辑
def create_couplet(api_key, base_url, model, upper_sentence, user_requirements):
    api = PoetryAPI(api_key, base_url, model, rpm_limit=10)
    result, status = api.couplet_stream(upper_sentence, user_requirements)
    return result, status


def create_poem(api_key, base_url, model, keywords, style, user_requirements):
    api = PoetryAPI(api_key, base_url, model, rpm_limit=10)
    result, status = api.generate_poem_stream(keywords, style, user_requirements)
    return result, status


# 定义 Gradio 界面组件
with gr.Blocks(title="儿童诗歌对句游戏") as demo:
    gr.Markdown("# ✨ 儿童诗歌对句游戏 ✨")

    with gr.Tab("诗句对句"):
        with gr.Row():
            with gr.Column(scale=1):
                api_key_couplet = gr.Textbox(label="API Key", type="password")
                base_url_couplet = gr.Textbox(label="Base URL", value="https://qianfan.baidubce.com/v2")
                model_couplet = gr.Dropdown(
                    label="模型",
                    choices=[
                        "ernie-bot-turbo",
                        "ernie-4.5-turbo-128k",
                        "ernie-4.5-turbo-vl-32k",
                        "deepseek-r1",
                        "ernie-speed-128k",
                        "llama-4-maverick-17b-128e-instruct"
                    ],
                    value="ernie-bot-turbo"
                )
                upper_sentence = gr.Textbox(label="诗句上半句", placeholder="例如：床前明月光")
                user_requirements_couplet = gr.Textbox(label="额外要求", placeholder="例如：解释再简单些")
                couplet_btn = gr.Button("生成对句")
            with gr.Column(scale=2):
                couplet_output = gr.Textbox(label="AI诗歌助手的对句", lines=15)
                couplet_status = gr.Textbox(label="状态", interactive=False)

    with gr.Tab("关键词创作"):
        with gr.Row():
            with gr.Column(scale=1):
                api_key_generate = gr.Textbox(label="API Key", type="password")
                base_url_generate = gr.Textbox(label="Base URL", value="https://qianfan.baidubce.com/v2")
                model_generate = gr.Dropdown(
                    label="模型",
                    choices=[
                        "ernie-bot-turbo",
                        "ernie-4.5-turbo-128k",
                        "ernie-4.5-turbo-vl-32k",
                        "deepseek-r1",
                        "ernie-speed-128k",
                        "llama-4-maverick-17b-128e-instruct"
                    ],
                    value="ernie-bot-turbo"
                )
                keywords = gr.Textbox(label="创作关键词", placeholder="例如：桃花、春天、小朋友")
                style = gr.Dropdown(label="诗歌风格", choices=["唐诗", "宋词", "五言绝句", "七言绝句"], value="唐诗")
                user_requirements_generate = gr.Textbox(label="额外要求", placeholder="例如：多写春天的景色、语言更活泼")
                generate_btn = gr.Button("生成诗歌")
            with gr.Column(scale=2):
                poem_output = gr.Textbox(label="AI诗歌助手的创作（含古诗推荐）", lines=15)
                poem_status = gr.Textbox(label="状态", interactive=False)

    # 绑定事件（已移除streaming参数）
    couplet_btn.click(
        fn=create_couplet,
        inputs=[api_key_couplet, base_url_couplet, model_couplet, upper_sentence, user_requirements_couplet],
        outputs=[couplet_output, couplet_status]
    )
    generate_btn.click(
        fn=create_poem,
        inputs=[api_key_generate, base_url_generate, model_generate, keywords, style, user_requirements_generate],
        outputs=[poem_output, poem_status]
    )

# 启动 Gradio 应用
if __name__ == "__main__":
    demo.launch()
